﻿/* Copyright 2007-2008 dnAnalytics Project.
 *
 * Contributors to this file:
 * Jurgen Van Gael
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 * 
 * * Redistributions of source code must retain the above copyright notice, this 
 *   list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright notice, 
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 * * Neither the name of the dnAnalytics Project nor the names of its contributors
 *   may be used to endorse or promote products derived from this software without
 *   specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
 * TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
using System;
using dnAnalytics.LinearAlgebra;
using dnAnalytics.LinearAlgebra.Decomposition;

namespace dnAnalytics.Statistics.Distributions
{
    /// <summary>
    /// This class implements functionality for matrix valued normal distributions. The distribution
    /// is parameterized by a mean matrix (M), a covariance matrix for the rows (V) and a covariance matrix
    /// for the columns (K). If the dimension of M is d-by-m then V is d-by-d and K is m-by-m.
    /// </summary>
    /// <remarks>The distribution will use the <see cref="System.Random"/> by default. 
    /// Users can set the random number generator by using the <see cref="RandomNumberGenerator"/> property.
    /// 
    /// The statistics classes will check all the incoming parameters whether they are in the allowed
    /// range. This might involve heavy computation. Optionally, by setting Control.CheckDistributionParameters
    /// to false, all parameter checks can be turned off.</remarks>
    public class MatrixNormal
    {
        // The mean of the matrix normal distribution.
        private readonly Matrix mM;

        // The covariance matrix for the rows.
        private readonly Matrix mV;

        // The covariance matrix for the columns.
        private readonly Matrix mK;

        /// <summary>
        /// Constructs a matrix normal distribution with specific mean and covariance matrices.
        /// </summary>
        /// <param name="M">The mean of the matrix normal.</param>
        /// <param name="V">The covariance matrix for the rows.</param>
        /// <param name="K">The covariance matrix for the columns.</param>
        /// <exception cref="ArgumentOutOfRangeException">If the dimensions of the mean and two covariance matrices don't match.</exception>
        public MatrixNormal(Matrix M, Matrix V, Matrix K)
        {
            if (Control.CheckDistributionParameters)
            {
                CheckParameters(M, V, K);
            }

            mM = M;
            mV = V;
            mK = K;

            RandomNumberGenerator = new System.Random();
        }

        /// <summary>
        /// A string representation of the distribution.
        /// </summary>
        public override string ToString()
        {
            return "MatrixNormal(Rows = " + mM.Rows + ", Columns = " + mM.Columns + ")";
        }

        private static void CheckParameters(Matrix M, Matrix V, Matrix K)
        {
            int d = M.Rows;
            int n = M.Columns;

            if (V.Columns != d || V.Rows != d)
            {
                throw new ArgumentOutOfRangeException("V", "V covariance matrix must have same number of rows as the mean matrix.");
            }
            if (K.Columns != n || K.Rows != n)
            {
                throw new ArgumentOutOfRangeException("K", "K covariance matrix must have same number of columns as the mean matrix.");
            }

            Cholesky cholV = new Cholesky(V);
            if (!cholV.IsPositiveDefinite())
            {
                throw new ArgumentOutOfRangeException("V", "A covariance matrix should be positive definite.");
            }

            Cholesky cholK = new Cholesky(K);
            if (!cholK.IsPositiveDefinite())
            {
                throw new ArgumentOutOfRangeException("K", "A covariance matrix should be positive definite.");
            }
        }

        /// <summary>
        /// The row dimensionality of the matrix normal distributed variable.
        /// </summary>
        public int Rows
        {
            get { return mM.Rows; }
        }

        /// <summary>
        /// The columns dimensionality of the matrix normal distributed variable.
        /// </summary>
        public int Columns
        {
            get { return mM.Columns; }
        }

        /// <summary>
        /// Gets the mean.
        /// </summary>
        /// <value>The mean.</value>
        public Matrix Mean
        {
            get { return mM; }
        }

        /// <summary>
        /// Gets the row covariance.
        /// </summary>
        /// <value>The row covariance.</value>
        public Matrix RowCovariance
        {
            get { return mV; }
        }

        /// <summary>
        /// Gets the column covariance.
        /// </summary>
        /// <value>The column covariance.</value>
        public Matrix ColumnCovariance
        {
            get { return mK; }
        }

        /// <summary>
        /// Samples a matrix normal distributed random variable.
        /// </summary>
        public Matrix Sample()
        {
            return Sample(RandomNumberGenerator, mM, mV, mK);
        }

        /// <summary>
        /// Samples a matrix normal distributed random variable.
        /// </summary>
        /// <param name="rnd">The random number generator to use.</param>
        /// <param name="M">The mean of the matrix normal.</param>
        /// <param name="V">The covariance matrix for the rows.</param>
        /// <param name="K">The covariance matrix for the columns.</param>
        /// <exception cref="ArgumentOutOfRangeException">If the dimensions of the mean and two covariance matrices don't match.</exception>
        public static Matrix Sample(System.Random rnd, Matrix M, Matrix V, Matrix K)
        {
            if (Control.CheckDistributionParameters)
            {
                CheckParameters(M, V, K);
            }

            int d = M.Rows;
            int n = M.Columns;

            // Compute the Kronecker product of V and K, this is the covariance matrix for the stacked matrix.
            Matrix VKi = V.KroneckerProduct(K.Inverse());

            // Sample a vector valued random variable with VKi as the covariance.
            Vector v = VectorNormal.Sample(rnd, new DenseVector(d*n, 0.0), VKi);

            // Unstack the vector v and add the mean.
            Matrix R = M.Clone();
            for (int i = 0; i < d; i++)
            {
                for (int j = 0; j < n; j++)
                {
                    R[i, j] += v[j*d + i];
                }
            }

            return R;
        }

        /// <summary>
        /// Gets or sets the random number generator.
        /// </summary>
        /// <value>The random number generator used to generate a random sample.</value>
        public System.Random RandomNumberGenerator { get; set; }

        /// <summary>
        /// Evaluates the probability density function for the matrix normal distribution.
        /// </summary>
        /// <param name="X">The matrix at which to evaluate the density at.</param>
        /// <exception cref="ArgumentOutOfRangeException">If the argument does not have the correct dimensions.</exception>
        public double Density(Matrix X)
        {
            if (X.Rows != mM.Rows || X.Columns != mM.Columns)
            {
                throw new ArgumentOutOfRangeException("X", "X must be of same size as the mean matrix.");
            }

            Matrix A = X - mM;
            Cholesky cholV = new Cholesky(mV);
            Cholesky cholK = new Cholesky(mK);

            return System.Math.Exp(-0.5*cholV.Solve(A.Transpose()*cholK.Solve(A)).Trace())
                   /System.Math.Pow(2.0*Math.Constants.Pi, X.Rows*X.Columns/2.0)
                   /System.Math.Pow(cholV.Determinant(), X.Rows/2.0)
                   /System.Math.Pow(cholK.Determinant(), X.Columns/2.0);
        }
    }
}